package com.gnote.langchain.service;

import com.gnote.langchain.domain.ChatMessages;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import lombok.RequiredArgsConstructor;
import org.bson.types.ObjectId;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.stereotype.Component;

import java.util.LinkedList;
import java.util.List;
import java.util.Objects;

/**
 * ai对话持久化存储
 *
 * @author Gnote
 * @since 2025/4/1 15:34
 */
@Component
@RequiredArgsConstructor
public class MongoChatMemoryStore implements ChatMemoryStore {
    private final MongoTemplate mongoTemplate;

    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        Criteria criteria = Criteria.where("memoryId").is(String.valueOf(memoryId));
        Query query = new Query(criteria);
        //更具id获取对话方式
        ChatMessages chatMessages = mongoTemplate.findOne(query, ChatMessages.class);
        if (chatMessages == null) return new LinkedList<>();
        //获取历史数据
        return ChatMessageDeserializer.messagesFromJson(chatMessages.getContent());
    }

    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        Query query = new Query(Criteria.where("memoryId").is(String.valueOf(memoryId)));
        //查询
        ChatMessages chatMessages = mongoTemplate.findOne(query, ChatMessages.class);
        if (Objects.isNull(chatMessages)) {
            //新增
            String content = ChatMessageSerializer.messagesToJson(messages);
            ChatMessages chatMessagesInsert = new ChatMessages(new ObjectId(), memoryId.toString(), content);
            mongoTemplate.insert(chatMessagesInsert);
        } else {
            //更新
            String content = ChatMessageSerializer.messagesToJson(messages);
            Update update = new Update();
            update.set("content", content);
            mongoTemplate.updateFirst(query, update, ChatMessages.class);
        }

    }

    @Override
    public void deleteMessages(Object memoryId) {
        //删除
        Query query = new Query(Criteria.where("memoryId").is(String.valueOf(memoryId)));
        mongoTemplate.remove(query, ChatMessages.class);
    }
}
